# -*- coding: utf-8 -*-

#    OriginDB
#    Copyright (C) 2010  Jason Power <power.jg@gmail.com>
#    Copyright (C) 2010  Travis Gockel <travis@gockelhut.com>
#
#    This program is free software; you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation; either version 2 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License along
#    with this program; if not, write to the Free Software Foundation, Inc.,
#    51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

from PyQt4 import QtCore
from PyQt4.QtGui import *

import sys
sys.path.append('../')
from aeDatabase import Pathway

from random import shuffle

from reactionVisual import ReactionVisual
from compoundVisual import CompoundVisual
from connection import Connection, Arrow

from tools import *


class PathwayScene(QGraphicsScene):

	def __init__(self, parent=None):
		QGraphicsScene.__init__(self, parent)

		self.pathway = Pathway()

		self.products = {}
		self.connections = []
		self.reactionElements = []

		self.oldSelection = []

		self.undoStack = QUndoStack()

		self.tool = ""

		self.reactStart = None
		self.cVisStart = None
		self.arrowItem = None

		self.connect(self, QtCore.SIGNAL("selectionChanged()"), self.newSelection)

	def setTool(self, tool):
		self.tool = tool

	def clear(self):
		self.products = {}
		self.connections = []
		self.reactionElements = []

		self.oldSelection = []

		self.undoStack.clear()

		QGraphicsScene.clear(self)

	def addItem(self, item):
		if type(item) == ReactionVisual:
			item.setReadOnly(False)
			self.reactionElements.append(item)
			
			# Add all of the products of this reaction to the list
			# so we can search it later (TODO do we need to update old things?)
			products = item.getProducts()
			for cVis in products:
				try:
					self.products[cVis.compound].append(item)
				except KeyError:
					self.products[cVis.compound] = [item]

			# For each reactant find if corresponding product
			reactants = item.getReactants()
			for rcVis in reactants:
				try:
					reactions = self.products[rcVis.compound]
				except KeyError:
					reactions = []

				# If there was a reaction found with same product then add connection
				for reaction in reactions:
					pcVis = reaction.getProduct(rcVis.compound)
					color = reaction.getColor(pcVis)
					if not color:
						shuffle(colors)
						color = colors[0]
						for color in colors:
							if reaction.testColor(color) and item.testColor(color):
								break
					
					conn = Connection(reaction, item, pcVis, rcVis, color)
					conn.add()
					self.connections.append(conn)
					QGraphicsScene.addItem(self, conn)
		
		if type(item) == Connection:
			self.connections.append(item)
		
		QGraphicsScene.addItem(self, item)


	def mousePressEvent(self, event):
		if event.button() == QtCore.Qt.LeftButton and self.tool == "erase":
			items = self.items(event.scenePos())
			item = None
			for i in items:
				if type(i) == ReactionVisual or type(i) == Connection:
					item = i
					break
			if not item:
				QGraphicsScene.mousePressEvent(self, event)
			if type(item) == ReactionVisual:
				command = RemoveReactionVisual(item, self)
				self.undoStack.push(command)
			elif type(item) == Connection:
				command = RemoveConnection(item, self)
				self.undoStack.push(command)
		
		elif event.button() == QtCore.Qt.LeftButton and self.tool == "draw":
			items = self.items(event.scenePos())
			item = None
			for i in items:
				if type(i) == ReactionVisual:
					item = i
					break
			if not item:
				QGraphicsScene.mousePressEvent(self, event)
			if type(item) == ReactionVisual:
				for i in items:
					if type(i) == CompoundVisual:
						self.cVisStart = i
						break
				if not self.cVisStart:
					return
				
				self.reactStart = item
				if self.cVisStart in self.reactStart.getProducts():
					self.arrowItem = Arrow()
					self.arrowItem.startPos = self.reactStart.getProductPosition(self.cVisStart)
					self.arrowItem.endPos = event.scenePos() + QtCore.QPointF(1,1)
					self.addItem(self.arrowItem)
					self.arrowItem.setZValue(100)
					event.accept()
					return
				else:
					self.reactStart = None
					self.cVisStart = None

		QGraphicsScene.mousePressEvent(self, event)


	def mouseMoveEvent(self, event):
		if self.arrowItem:
			self.arrowItem.endPos = event.scenePos()
			self.arrowItem.calcGeometry()
			self.arrowItem.update()
		
		QGraphicsScene.mouseMoveEvent(self, event)


	def mouseReleaseEvent(self, event):
		if self.arrowItem:
			self.removeItem(self.arrowItem)
			self.arrowItem = None
		
		if not self.reactStart:
			QGraphicsScene.mouseReleaseEvent(self, event)
			return
		
		items = self.items(event.scenePos())
		item = None
		for i in items:
			if type(i) == ReactionVisual or type(i) == Connection:
				item = i
				break
		if not item:
			QGraphicsScene.mousePressEvent(self, event)
		
		if type(item) != ReactionVisual or item == self.reactStart:
			QGraphicsScene.mouseReleaseEvent(self, event)
			return

		cVisEnd = None
		for i in items:
			if type(i) == CompoundVisual:
				cVisEnd = i
				break

		if not cVisEnd in item.getReactants():
			return

		if item == self.reactStart or (cVisEnd.compound != self.cVisStart.compound):
			return
		
		for conn in self.connections:
			if conn.reactantReaction == item and conn.productReaction == self.reactStart:
				return
		
		connection = Connection(self.reactStart, item, self.cVisStart, cVisEnd)
		self.connections.append(connection)
		command = AddConnection(connection, self)
		self.undoStack.push(command)

		self.reactStart = None
		self.cVisStart = None


	def newSelection(self):
		for item in self.oldSelection:
			item.setSelected(False)
		self.oldSelection = self.selectedItems()
		for item in self.oldSelection:
			item.setSelected(True)

	def getPathway(self):
		self.pathway = Pathway()
		graph = []
		for conn in self.connections:
			graph.append(conn.getSerial())
		print graph
		self.pathway.graph = graph

		return self.pathway
	
	def setPathway(self, pathway):
		self.pathway = pathway
		
		for edge in pathway:
			productReaction = edge[0]
			reactantReaction = edge[1][0]
			compound = edge[1][1]

			reactantVis = None
			productVis = None
			for rvis in self.reactionElements:
				if rvis.reaction == reactantReaction:
					reactantVis = rvis
				if rvis.reaction == productReaction:
					productVis = rvis
			
			if not reactantVis:
				reactantVis = ReactionVisual(reactantReaction)
				self.reactionElements.append(reactantVis)
				QGraphicsScene.addItem(self, reactantVis)
				for cVis in reactantVis.getProducts():
					try:
						self.products[cVis.compound].append(reactantVis)
					except KeyError:
						self.products[cVis.compound] = [reactantVis]

			if not productVis:
				productVis = ReactionVisual(productReaction)
				self.reactionElements.append(productVis)
				QGraphicsScene.addItem(self, productVis)
				for cVis in productVis.getProducts():
					try:
						self.products[cVis.compound].append(productVis)
					except KeyError:
						self.products[cVis.compound] = [productVis]

			pcVis = productVis.getProduct(compound)
			rcVis = reactantVis.getReactant(compound)

			color = productVis.getColor(pcVis)
			if not color:
				color = reactantVis.getColor(rcVis)
			
			if not color:
				shuffle(colors)
				color = colors[0]
				for color in colors:
					if productVis.testColor(color) and reactantVis.testColor(color):
						break

			conn = Connection(productVis, reactantVis, pcVis, rcVis, color)
			conn.add()
			self.connections.append(conn)
			QGraphicsScene.addItem(self, conn)

		self.clean()

	def removeItem(self, item):
		if type(item) == ReactionVisual:
			toRemove = []
			i = 0
			for conn in self.connections:
				if conn.productReaction == item or conn.reactantReaction == item:
					QGraphicsScene.removeItem(self, conn)
					conn.remove()
					toRemove.append(i)
				i += 1
			for i in range(len(toRemove)):
				del self.connections[toRemove[i]-i]
			for cVis in item.getProducts():
				for i in range(len(self.products[cVis.compound])):
					if self.products[cVis.compound][i] == item:
						del self.products[cVis.compound][i]
						break
			for i in range(len(self.reactionElements)):
				if self.reactionElements[i] == item:
					del self.reactionElements[i]
					break
			
		if type(item) == Connection:
			for i in range(len(self.connections)):
				if self.connections[i] == item:
					del self.connections[i]
					break
		
		QGraphicsScene.removeItem(self, item)
		self.update()

	def clean(self):

		def findNext(reaction):
			ans = []
			for i in self.connections:
				if i.productReaction == reaction:
					ans.append(i.reactantReaction)
			return ans

		# get the zeroth order elements
		lastOrder = set()
		for r in self.reactionElements:
			add = True
			for conn in self.connections:
				if conn.reactantReaction == r:
					add = False
					break
			if add:
				lastOrder.add(r)

		totalY = 0
		for rVis in lastOrder:
			totalY += rVis.boundingRect().height()

		totalY += 20 * (len(lastOrder) - 1)

		#print "first order", lastOrder

		currentX = 0
		maxX = 0
		y = -totalY/2
		for r in lastOrder:
			r.setPos(currentX, y)
			y += r.boundingRect().height()
			y += 20
			maxX = max(maxX, r.boundingRect().width())

		currentX += maxX + 50
		
		while lastOrder:
			nextOrder = set()
			for r in lastOrder:
				nextOrder.update(findNext(r))

			toRemove = set()
			for ri in nextOrder:
				children = findNext(ri)
				for rj in nextOrder:
					if ri==rj:
						continue
					if rj in children:
						toRemove.add(rj)

			for r in toRemove:
				nextOrder.remove(r)

			totalY = 0
			for rVis in nextOrder:
				totalY += rVis.boundingRect().height()

			totalY += 20 * (len(nextOrder) - 1)

			maxX = 0
			y = -totalY/2
			for r in nextOrder:
				r.setPos(currentX, y)
				y += r.boundingRect().height()
				y += 20
				maxX = max(maxX, r.boundingRect().width())

			currentX += maxX + 50

			#print nextOrder

			lastOrder = nextOrder

		self.update()
